import time

from fastapi.middleware.gzip import GZipMiddleware
from fastapi import FastAPI, Request
from fastapi.openapi.models import Response
from starlette.middleware.base import BaseHTTPMiddleware
from datetime import datetime


def register_json_result(app: FastAPI):
    # 将中间件添加到主程序中
    app.add_middleware(ProcessTimeMiddleware, header_namespace="middleware")


# 将时间格式化为字符串
def _time2str(time_str):
    return datetime.strftime(time_str, '%Y-%m-%d %H:%M:%S')


# 将字符串转换为时间
def _str2time(time_str):
    return datetime.strptime(time_str, '%Y-%m-%d %H:%M:%S')


class ProcessTimeMiddleware(BaseHTTPMiddleware):
    def __init__(self, app, header_namespace: str):
        super().__init__(app)
        # 自定义参数，用于定义middleware的header名称空间
        self.header_namespace = header_namespace

    async def dispatch(self, request: Request, call_next):
        # 接收来自客户端的Request请求；
        headers = dict(request.scope['headers'])
        # 定义middleware开始时间
        middleware_start_time = _time2str(datetime.now())
        # 将middleware开始时间添加到request的headers中，这里request.headers是一个可读可写的对象，但是它的值是不可变的，所以这里需要将request.headers转换为字典，然后再修改字典的值，最后再将字典转换为元组，赋值给request.scope['headers']；
        headers[b'middleware_start_time'] = middleware_start_time.encode('utf-8')
        request.scope['headers'] = [(k, v) for k, v in headers.items()]

        # 将Request请求传回原路由
        response: Response = await call_next(request)

        # 接收来自原路由的Response响应，将middleware结束时间添加到response的headers中
        response.headers[f"{self.header_namespace}_start_time"] = middleware_start_time
        response.headers[f"{self.header_namespace}_end_time"] = _time2str(datetime.now())

        return response
